{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d6509c3a",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/callbacks/WandbCallbackHandler.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0d8b66c",
   "metadata": {},
   "source": [
    "# Wandb Callback Handler\n",
    "\n",
    "[Weights & Biases Prompts](https://docs.wandb.ai/guides/prompts) is a suite of LLMOps tools built for the development of LLM-powered applications.\n",
    "\n",
    "The `WandbCallbackHandler` is integrated with W&B Prompts to visualize and inspect the execution flow of your index construction, or querying over your index and more. You can use this handler to persist your created indices as W&B Artifacts allowing you to version control your indices.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c3527e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-callbacks-wandb\n",
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "612f35ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OpenAI API key configured\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "if os.getenv(\"OPENAI_API_KEY\") is None:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\n",
    "        \"Paste your OpenAI key from:\"\n",
    "        \" https://platform.openai.com/account/api-keys\\n\"\n",
    "    )\n",
    "assert os.getenv(\"OPENAI_API_KEY\", \"\").startswith(\n",
    "    \"sk-\"\n",
    "), \"This doesn't look like a valid OpenAI API key\"\n",
    "print(\"OpenAI API key configured\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78a29d9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.callbacks import CallbackManager\n",
    "from llama_index.core.callbacks import LlamaDebugHandler\n",
    "from llama_index.callbacks.wandb import WandbCallbackHandler\n",
    "from llama_index.core import (\n",
    "    VectorStoreIndex,\n",
    "    SimpleDirectoryReader,\n",
    "    SimpleKeywordTableIndex,\n",
    "    StorageContext,\n",
    ")\n",
    "from llama_index.llms.openai import OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6feb252",
   "metadata": {},
   "source": [
    "## Setup LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d22fee33",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = OpenAI(model=\"gpt-4\", temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8790f4c7",
   "metadata": {},
   "source": [
    "## W&B Callback Manager Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a32b984-772e-4832-945e-cb6fc7be9e0b",
   "metadata": {},
   "source": [
    "**Option 1**: Set Global Evaluation Handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a3b9d22-cd67-4fb5-9785-254e58179a02",
   "metadata": {},
   "outputs": [],
   "source": [
    "import llama_index.core\n",
    "from llama_index.core import set_global_handler\n",
    "\n",
    "set_global_handler(\"wandb\", run_args={\"project\": \"llamaindex\"})\n",
    "wandb_callback = llama_index.core.global_handler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1755516-f8ad-458e-b52f-f7665c023e43",
   "metadata": {},
   "source": [
    "**Option 2**: Manually Configure Callback Handler\n",
    "\n",
    "Also configure a debugger handler for extra notebook visibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "defa9155-daca-4a8f-8ca6-87d1ee98f084",
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_debug = LlamaDebugHandler(print_trace_on_end=True)\n",
    "\n",
    "# wandb.init args\n",
    "run_args = dict(\n",
    "    project=\"llamaindex\",\n",
    ")\n",
    "\n",
    "wandb_callback = WandbCallbackHandler(run_args=run_args)\n",
    "\n",
    "Settings.callback_manager = CallbackManager([llama_debug, wandb_callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4cf969a",
   "metadata": {},
   "source": [
    "> After running the above cell, you will get the W&B run page URL. Here you will find a trace table with all the events tracked using [Weights and Biases' Prompts](https://docs.wandb.ai/guides/prompts) feature."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4a7c101",
   "metadata": {},
   "source": [
    "## 1. Indexing"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e5d31f80",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e7ad71e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1011596",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d6975c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********\n",
      "Trace: index_construction\n",
      "    |_node_parsing ->  0.295179 seconds\n",
      "      |_chunking ->  0.293976 seconds\n",
      "    |_embedding ->  0.494492 seconds\n",
      "    |_embedding ->  0.346162 seconds\n",
      "**********\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Logged trace tree to W&B.\n"
     ]
    }
   ],
   "source": [
    "index = VectorStoreIndex.from_documents(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a948efc",
   "metadata": {},
   "source": [
    "### 1.1 Persist Index as W&B Artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ad58e67",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Adding directory to artifact (/Users/loganmarkewich/llama_index/docs/examples/callbacks/wandb/run-20230801_152955-ds93prxa/files/storage)... Done. 0.0s\n"
     ]
    }
   ],
   "source": [
    "wandb_callback.persist_index(index, index_name=\"simple_vector_store\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ed156a6",
   "metadata": {},
   "source": [
    "### 1.2 Download Index from W&B Artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc35f448",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m:   3 of 3 files downloaded.  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********\n",
      "Trace: index_construction\n",
      "**********\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core import load_index_from_storage\n",
    "\n",
    "storage_context = wandb_callback.load_storage_context(\n",
    "    artifact_url=\"ayut/llamaindex/simple_vector_store:v0\"\n",
    ")\n",
    "\n",
    "# Load the index and initialize a query engine\n",
    "index = load_index_from_storage(\n",
    "    storage_context,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae4de4a9",
   "metadata": {},
   "source": [
    "## 2. Query Over Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42221465",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**********\n",
      "Trace: query\n",
      "    |_query ->  2.695958 seconds\n",
      "      |_retrieve ->  0.806379 seconds\n",
      "        |_embedding ->  0.802871 seconds\n",
      "      |_synthesize ->  1.8893 seconds\n",
      "        |_llm ->  1.842434 seconds\n",
      "**********\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Logged trace tree to W&B.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The text does not provide information on what the author did growing up.\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"What did the author do growing up?\")\n",
    "print(response, sep=\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c49ff101",
   "metadata": {},
   "source": [
    "## Close W&B Callback Handler\n",
    "\n",
    "When we are done tracking our events we can close the wandb run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28ef6a7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb_callback.finish()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
